{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Template Attack with Hardware Assumption\n",
    "\n",
    "Supported setups:\n",
    "\n",
    "SCOPES:\n",
    "\n",
    "* OPENADC\n",
    "* CWNANO\n",
    "\n",
    "PLATFORMS:\n",
    "\n",
    "* CWLITEARM\n",
    "* CWLITEXMEGA\n",
    "* CWNANO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Template attacks* are a powerful type of side-channel attack. These attacks are a subset of *profiling attacks*, where an attacker creates a \"profile\" of a sensitive device and applies this profile to quickly find a victim's secret key. \n",
    "\n",
    "Template attacks require more setup than CPA attacks. To perform a template attack, the attacker must have access to another copy of the protected device that they can fully control. Then, they must perform a great deal of pre-processing to create the template - in practice, this may take dozens of thousands of power traces. However, the advantages are that template attacks require a very small number of traces from the victim to complete the attack. With enough pre-processing, the key may be able to be recovered from just a single trace. \n",
    "\n",
    "There are four steps to a template attack:\n",
    "1. Using a copy of the protected device, record a large number of power traces using many different inputs (plaintexts and keys). Ensure that enough traces are recorded to give us information about each subkey value.\n",
    "2. Create a template of the device's operation. This template notes a few \"points of interest\" in the power traces and a multivariate distribution of the power traces at each point. \n",
    "3. On the victim device, record a small number of power traces. Use multiple plaintexts. (We have no control over the secret key, which is fixed.)\n",
    "4. Apply the template to the attack traces. For each subkey, track which value is most likely to be the correct subkey. Continue until the key has been recovered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "num_traces = 50\n",
    "CHECK_CORR = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../hardware/victims/firmware/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Signals, Noise, and Statistics\n",
    "Before looking at the details of the template attack, it is important to understand the statistics concepts that are involved. A template is effectively a multivariate distribution that describes several key samples in the power traces. This section will describe what a multivariate distribution is and how it can be used in this context."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Noise Distributions\n",
    "Electrical signals are inherently noisy. Any time we take a voltage measurement, we don't expect to see a perfect, constant level. For example, if we attached a multimeter to a 5 V source and took 4 measurements, we might expect to see a data set like (4.95, 5.01, 5.06, 4.98). One way of modelling this voltage source is:\n",
    "\n",
    "$$\\mathbf{X} = X_{actual} + \\mathbf{N}$$\n",
    "\n",
    "A simple model for these random variables uses a Gaussian distribution (read: a bell curve). The probability density function (PDF) of a Gaussian distribution is\n",
    "\n",
    "$$f(x) = \\frac{1}{\\sigma \\sqrt{2\\pi}} e^{-(x - \\mu)^2 / 2\\sigma^2}$$\n",
    "\n",
    "where $\\mu$ is the mean and $\\sigma$ is the standard deviation. For instance, our voltage source might have a mean of 5 and a standard deviation of 0.5. We can plot this curve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 5\n",
    "sigma = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's this graph all about? This graph shows the *probability density* of this random variable. Note it does NOT give you an actual probability, instead you need to integrate to find the probability the variable is lieing in some range. This is important because the density function won't always return just < 1, and with multivariate (to be discussed) the output is frequently larger than one. You can find probabilities using the CDF, which integrates from negative infinity to the given random variable value.\n",
    "\n",
    "For example, what is the chance that we will observe a value in the range 6.0?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import stats\n",
    "pval = stats.norm.cdf(6.0, mu, sigma)\n",
    "print(\"%.2f %%\"%(pval*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PDFs help us answer a simple question (that will be relevant to us shortly). Let's say we have the following experiment setup:\n",
    "\n",
    "\n",
    "We have two batteries and are trying to figure out which one is connected behind the curtain. There is a *lot* of noise in our measurement setup and we don't have much time to waste, so can only measure a few points. The following will simulate a number of measurements from one or the other battery, but not tell us which one is which."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import randint\n",
    "import numpy as np\n",
    "\n",
    "noise = 2.0\n",
    "\n",
    "def do_measure():\n",
    "    if randint(0,1):\n",
    "        return (1.5, np.random.normal(1.5, noise))\n",
    "    else:\n",
    "        return (3.0, np.random.normal(3.0, noise))\n",
    " \n",
    "N = 10\n",
    "\n",
    "measure_array = [0]*N\n",
    "answer = [0]*N\n",
    "\n",
    "for n in range(0, N):\n",
    "    answer[n], measure_array[n] = do_measure()\n",
    "    print(\"%.02f\"%measure_array[n])\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how you stack up to a computer! Try entering your guess of the measurements here by editing the values of the array to reflect the printed measurements. Select if you think the source was 1.5V or 3.0V."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "guess = [1.5, 1.5, 3.0, 3.0, 1.5, 1.5, 3.0, 1.5, 3.0, 3.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How does the PDF and CDF information help us? We could use the PDF to provide some information on the liklihood of a specific input. Note this should be done with another function (remember the PDF is a density function, not a probability function), but we will use PDF for now as the template attack to be discussed will continue this method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "results = []*N\n",
    "\n",
    "for n in range(0, N):\n",
    "    p15 = stats.norm.pdf(measure_array[n], 1.5, noise)\n",
    "    p30 = stats.norm.pdf(measure_array[n], 3.0, noise)\n",
    "\n",
    "    if p15 > p30:\n",
    "        bguess = 1.5\n",
    "    else:\n",
    "        bguess = 3.0\n",
    "        \n",
    "    results.append([p15, p30, bguess, guess[n], answer[n]])\n",
    "    \n",
    "pd.DataFrame(results, columns=[\"P|m=1.5\", \"P|m=3.0\", \"Best Guess\", \"Your Guess\", \"Answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multivariate Statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 1-variable Gaussian distribution works well for one measurement. What if we're working with more than one random variable?\n",
    "\n",
    "Suppose we're measuring two voltages that have some amount of noise on them. We'll call them $\\mathbf{X}$ and $\\mathbf{Y}$. As a first attempt, we could write down a model for $\\mathbf{X}$ using a normal distribution and a separate model for $\\mathbf{Y}$ using a different distribution. However, this might not always make sense. If we write two separate distributions, what we're saying is that the two variables are independent: when $\\mathbf{X}$ goes up, there's no guarantee that $\\mathbf{Y}$ will follow it.\n",
    "\n",
    "Multivariate distributions let us model multiple random variables that may or may not be correlated. In a multivariate distribution, instead of writing down a single variance $\\sigma$, we keep track of a whole matrix of covariances. For example, to model three random variables ($\\mathbf{X}, \\mathbf{Y}, \\mathbf{Z}$), this matrix would be\n",
    "\n",
    "$$\n",
    "\\mathbf{\\Sigma} = \n",
    "\\begin{bmatrix}\n",
    "Var(\\mathbf{X})             & Cov(\\mathbf{X}, \\mathbf{Y}) & Cov(\\mathbf{X}, \\mathbf{Z}) \\\\\n",
    "Cov(\\mathbf{Y}, \\mathbf{X}) & Var(\\mathbf{Y})             & Cov(\\mathbf{Y}, \\mathbf{Z}) \\\\\n",
    "Cov(\\mathbf{Z}, \\mathbf{X}) & Cov(\\mathbf{Z}, \\mathbf{Y}) & Var(\\mathbf{Z}) \n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "Also, note that this distribution needs to have a mean for each random variable:\n",
    "\n",
    "$$\n",
    "\\mathbf{\\mu} = \n",
    "\\begin{bmatrix}\n",
    "\\mu_X \\\\\n",
    "\\mu_Y \\\\\n",
    "\\mu_Z\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "The PDF of this distribution is more complicated: instead of using a single number as an argument, it uses a vector with all of the variables in it ($\\mathbf{x} = [x, y, z, \\dots]^T$). The equation for $k$ random variables is\n",
    "\n",
    "$$\n",
    "f(\\mathbf{x})\n",
    "= \\frac{1}{\\sqrt{(2\\pi)^k |\\mathbf{\\Sigma}|}} \n",
    "  e^{-(\\mathbf{(x - \\mu)}^T \\mathbf{\\Sigma}^{-1} \\mathbf{(x - \\mu)} / 2}\n",
    "$$\n",
    "\n",
    "Don't worry if this looks crazy - the SciPy package in Python will do all the heavy lifting for us. As with the single-variable distributions, we're going to use this to find how likely a certain observation is. In other words, if we put $k$ points of our power trace into $\\mathbf{x}$ and we find that $f(\\mathbf{x})$ is very high, then we've probably found a good guess.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the Template\n",
    "\n",
    "A template is a set of probability distributions that describe what the\n",
    "power traces look like for many different keys. Effectively, a template\n",
    "says: \\\"If you\\'re going to use key $k$, your power trace will look like\n",
    "the distribution $f_k(\\mathbf{x})$\\\". We can use this information to\n",
    "find subtle differences between power traces and to make very good key\n",
    "guesses for a single power trace.\n",
    "\n",
    "### Number of Traces\n",
    "\n",
    "One of the downsides of template attacks is that they require a great\n",
    "number of traces to be preprocessed before the attack can begin. This is\n",
    "mainly for statistical reasons. In order to come up with a good\n",
    "distribution to model the power traces for *every key*, we need a large\n",
    "number of traces for *every key*. For example, if we\\'re going to attack\n",
    "a single subkey of AES-128, then we need to create 256 power consumption\n",
    "models (one for every number from 0 to 255). In order to get enough data\n",
    "to make good models, we need tens of thousands of traces.\n",
    "\n",
    "Note that we don\\'t have to model every single key. One good alternative\n",
    "is to model a sensitive part of the algorithm, like the substitution box\n",
    "in AES. We can get away with a much smaller number of traces here; if we\n",
    "make a model for every possible Hamming weight, then we would end up\n",
    "with 9 models, which is an order of magnitude smaller. However, then we\n",
    "can\\'t recover the key from a single attack trace - we need more\n",
    "information to recover the secret key.\n",
    "\n",
    "### Points of Interest\n",
    "\n",
    "Our goal is to create a multivariate probability describing the power\n",
    "traces for every possible key. If we modeled the entire power trace this\n",
    "way (with, say, 3000 samples), then we would need a 3000-dimension\n",
    "distribution. This is insane, so we\\'ll find an alternative.\n",
    "\n",
    "Thankfully, not every point on the power trace is important to us. There\n",
    "are two main reasons for this:\n",
    "\n",
    "-   We might be taking more than one sample per clock cycle. (Through\n",
    "    most of the ChipWhisperer tutorials, our ADC runs four times faster\n",
    "    than the target device.) There\\'s no real reason to use all of these\n",
    "    samples - we can get just as much information from a single sample\n",
    "    at the right time.\n",
    "-   Our choice of key doesn\\'t affect the entire power trace. It\\'s\n",
    "    likely that the subkeys only influence the power consumption at a\n",
    "    few critical times. If we can pick these important times, then we\n",
    "    can ignore most of the samples.\n",
    "\n",
    "These two points mean that we can usually live with a handful (3-5) of\n",
    "*points of interest*. If we can pick out good points and write down a\n",
    "model using these samples, then we can use a 3D or 5D distribution - a\n",
    "great improvement over the original 3000D model.\n",
    "\n",
    "#### Picking POIs\n",
    "\n",
    "There are several ways to pick the most important points in each of the\n",
    "traces. Generally, the aim is to find points that vary strongly between\n",
    "different operations (subkeys or Hamming weights). The simplest method\n",
    "\\-- the one that we\\'ll use here \\-- is the *sum of differences* method.\n",
    "\n",
    "The algorithm for the sum of difference method is:\n",
    "\n",
    "-   For every operation $k$ and every sample $i$, find the average power\n",
    "    $M_{k, i}$. For instance, if there are $T_k$ traces where we\n",
    "    performed operation $k$, then this average is\n",
    "\n",
    "$M_{k, i} = \\frac{1}{T_k} \\sum_{j=1}^{T_k} t_{j, i}$\n",
    "\n",
    "-   After finding all of the means, calculate all of their absolute\n",
    "    pairwise differences. Add these up. This will give one \\\"trace\\\"\n",
    "    which has peaks where the samples are usually different. The\n",
    "    calculation looks like\n",
    "\n",
    "$D_{i} = \\sum_{k_1, k_2} |M_{k_1, i} - M_{k_2, i}|$\n",
    "\n",
    "An example of this sum of differences is:\n",
    "\n",
    "![][1]\n",
    "\n",
    "-   The peaks of $D_i$ show the most important points, but we need to\n",
    "    satisfy point 1 from above - we need to pick some peaks that aren\\'t\n",
    "    too close. One algorithm to do this is:\n",
    "\n",
    "1.  Pick the highest point in $D_i$ and save this value of $i$ as a\n",
    "    point of interest. (ie: $i = argmax(D_i)$)\n",
    "2.  Throw out the nearest $N$ points (where $N$ is the minimum spacing\n",
    "    between POIs).\n",
    "3.  Repeat until enough POIs have been selected.\n",
    "\n",
    "### Analyzing the Data\n",
    "\n",
    "Suppose that we\\'ve picked $I$ points of interest, which are at samples\n",
    "$s_i$ ($0 \\le i < I$). Then, our goal is to find a mean and covariance\n",
    "matrix for every operation (every choice of subkey or intermediate\n",
    "Hamming weight). Let\\'s say that there are $K$ of these operations\n",
    "(maybe 256 subkeys or 9 possible Hamming weights).\n",
    "\n",
    "For now, we\\'ll look at a single operation $k$ ($0 \\le k < K$). The\n",
    "steps are:\n",
    "\n",
    "-   Find every power trace $t$ that falls under the category of\n",
    "    \\\"operation $k$\\\". (ex: find every power trace where we used a\n",
    "    subkey of 0x01.) We\\'ll say that there are $T_k$ of these, so\n",
    "    $t_{j, s_i}$ means the value at trace $j$ and POI $i$.\n",
    "-   Find the average power $\\mu_i$ at every point of interest. This\n",
    "    calculation will look like:\n",
    "\n",
    "$\\mu_i = \\frac{1}{T_k} \\sum_{j=1}^{T_k} t_{j, s_i}$\n",
    "\n",
    "-   Find the variance $v_i$ of the power at each point of interest. One\n",
    "    way of calculating this is:\n",
    "\n",
    "$v_i = \\frac{1}{T_k} \\sum_{j=1}^{T_k} (t_{j, s_i} - \\mu_i)^2$\n",
    "\n",
    "-   Find the covariance $c_{i, i^*}$ between the power at every pair of\n",
    "    POIs ($i$ and $i^*$). One way of calculating this is:\n",
    "\n",
    "$c_{i, i^*} = \\frac{1}{T_k}  \\sum_{j=1}^{T_k} (t_{j, s_i} - \\mu_i) (t_{j, s_{i^*}} - \\mu_{i^*})$\n",
    "\n",
    "-   Put together the mean and covariance matrices as:\n",
    "\n",
    "$\\mathbf{\\mu} = \n",
    "\\begin{bmatrix}\n",
    "\\mu_1 \\\\\n",
    "\\mu_2 \\\\\n",
    "\\mu_3 \\\\\n",
    "\\vdots\n",
    "\\end{bmatrix}$\n",
    "\n",
    "$\\mathbf{\\Sigma} = \n",
    "\\begin{bmatrix}\n",
    "v_1     & c_{1,2} & c_{1,3} & \\dots \\\\\n",
    "c_{2,1} & v_2     & c_{2,3} & \\dots \\\\\n",
    "c_{3,1} & c_{3,2} & v_3     & \\dots \\\\\n",
    "\\vdots  & \\vdots  & \\vdots  & \\ddots \n",
    "\\end{bmatrix}$\n",
    "\n",
    "These steps must be done for every operation $k$. At the end of this\n",
    "preprocessing, we\\'ll have $K$ mean and covariance matrices, modelling\n",
    "each of the $K$ different operations that the target can do.\n",
    "\n",
    "Using the Template\n",
    "------------------\n",
    "\n",
    "With a template in hand, we can finish our attack. For the attack, we\n",
    "need a smaller number of traces - we\\'ll say that we have $A$ traces.\n",
    "The sample values will be labeled $a_{j, s_i}$ ($1 \\le j \\le A$).\n",
    "\n",
    "### Applying the Template\n",
    "\n",
    "First, let\\'s apply the template to a single trace. Our job is to decide\n",
    "how likely all of our key guesses are. We need to do the following:\n",
    "\n",
    "-   Put our trace values at the POIs into a vector. This vector will be\n",
    "\n",
    "$\\mathbf{a_j} = \n",
    "\\begin{bmatrix}\n",
    "a_{j,1} \\\\\n",
    "a_{j,2} \\\\\n",
    "a_{j,3} \\\\\n",
    "\\vdots\n",
    "\\end{bmatrix}$\n",
    "\n",
    "-   Calculate the PDF for every key guess and save these for later. This\n",
    "    might look like:\n",
    "\n",
    "$p_{k, j} = f_k(\\mathbf{a_j})$\n",
    "\n",
    "-   Repeat these two steps for all of the attack traces.\n",
    "\n",
    "This process gives us an array of $p_{k, j}$, which says: \\\"Looking at\n",
    "trace $j$, how likely is it that key $k$ is the correct one?\\\"\n",
    "\n",
    "### Combining the Results\n",
    "\n",
    "The very last step is to combine our $p_{k, j}$ values to decide which\n",
    "key is the best fit. The easiest way to do this is to combine them as\n",
    "\n",
    "$P_k = \\prod_{j=1}^{A} p_{k,j}$\n",
    "\n",
    "For example, if we guessed that a subkey was equal to 0x00 and our PDF\n",
    "results in 3 traces were (0.9, 0.8, 0.95), then our overall result would\n",
    "be 0.684. Having one trace that doesn\\'t match the template can cause\n",
    "this number to drop quickly, helping us eliminate the wrong guesses.\n",
    "Finally, we can pick the highest value of $P_k$, which tells us which\n",
    "guess fits the templates the best, and we\\'re done!\n",
    "\n",
    "This method of combining our per-trace results can suffer from precision\n",
    "issues. After multiplying many large or small numbers together, we could\n",
    "end up with numbers that are too large or small to fit into a floating\n",
    "point variable. An easy fix is to work with logarithms. Instead of using\n",
    "$P_k$ directly, we can calculate\n",
    "\n",
    "$\\log P_k = \\sum_{j=1}^A \\log p_{k,j}$\n",
    "\n",
    "Comparing these logarithms will give us the same results without the\n",
    "precision issues.\n",
    "\n",
    "![1](https://wiki.newae.com/images/9/97/Template-Sum-Of-Difference.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capturing Power Traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = \"../hardware/victims/firmware/simpleserial-aes/simpleserial-aes-{}.hex\".format(PLATFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# program the target\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Capturing Template (Rand Key, Rand Text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first capture will be for the template. This template requires us to have full control of the device, both the key and plaintext. We will be using a **random key**, meaning the key will change on **every** encryption.\n",
    "\n",
    "Why would we do that? The reason is that if we used a **fixed key** during the profile, there is a strong chance that we will \"overfit\" the templates. That is the templates will only be valid for the fixed key we used during the profiling phase. Obviously this won't make for a useful attack in practice, since we don't know the actual key in use.\n",
    "\n",
    "Our \"overfit\" model would also mean our results appear to work increadily well, but they won't work in a real situation. That is, should you perform profiling and evaluation with the same fixed key, you will get very good results due to the overfit but they are meaningless in practice.\n",
    "\n",
    "The capture phase is otherwise similar to before. The difference here is that we will have to tell the acquisition model that we want to use a random key. The following simple block will create a new project and save 2000 traces.\n",
    "\n",
    "Remember that the resulting group of Hamming weights won't be uniformally distributed. This will affect our template results, as for example the group with a Hamming weight of 8 will have only about $1/256 * 2000 = 8$ traces. This is because the Hamming weight of 8 only exists for 1/256 possible values of the S-Box output. Hamming weights of 3,4,5 are much more likely for example and will contain much larger trace groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Capture Traces\n",
    "from tqdm import tnrange\n",
    "import chipwhisperer as cw\n",
    "\n",
    "#Make a project to save template traces\n",
    "project = cw.create_project(\"projects/tutorial_template_templatedata.cwp\", overwrite = True)\n",
    "\n",
    "ktp = cw.ktp.Basic()\n",
    "ktp.fixed_key = False #RANDOM KEY in addition to RANDOM TEXT\n",
    "\n",
    "N = 6000  # Number of traces\n",
    "for i in tnrange(N, desc='Capturing traces'):\n",
    "    key, text = ktp.next()\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)\n",
    "\n",
    "#Save the project\n",
    "project.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While we're here - why not capture the traces for the attack too. These traces will be much shorter, we'll only capture 20 traces and probably not even need all of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Capture Traces\n",
    "from tqdm import tnrange\n",
    "import chipwhisperer as cw\n",
    "import numpy as np\n",
    "\n",
    "#Make a project to save test traces\n",
    "project = cw.create_project(\"projects/tutorial_template_validate.cwp\", overwrite = True)\n",
    "\n",
    "ktp = cw.ktp.Basic()\n",
    "\n",
    "N = 20  # Number of traces\n",
    "for i in tnrange(N, desc='Capturing traces'):\n",
    "    key, text = ktp.next()\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)\n",
    "\n",
    "#Save the project\n",
    "project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup the connection to the target and scope\n",
    "scope.dis()\n",
    "target.dis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Opening & Grouping Traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "project_template = cw.open_project(\"projects/tutorial_template_templatedata.cwp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Let's confirm that we get random keys\n",
    "for i in range(0, 4):\n",
    "    print(\"%d: \"%i + \" \".join([\"%02x\"%project_template.keys[i][j] for j in range(0, 16)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "hv.extension('bokeh')\n",
    "hv.Curve(project_template.waves[0]).opts(height=600, width=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "intermediate = sbox[plaintext ^ key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbox=(\n",
    "    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,\n",
    "    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,\n",
    "    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,\n",
    "    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,\n",
    "    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,\n",
    "    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,\n",
    "    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,\n",
    "    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,\n",
    "    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,\n",
    "    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,\n",
    "    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,\n",
    "    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,\n",
    "    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,\n",
    "    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,\n",
    "    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,\n",
    "    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16) \n",
    "\n",
    "hw = [bin(x).count(\"1\") for x in range(256)]\n",
    "\n",
    "for n in range(0, 10):\n",
    "    tin = project_template.textins[n][0]\n",
    "    key = project_template.keys[n][0]\n",
    "    \n",
    "    intermediate = sbox[tin ^ key]\n",
    "    \n",
    "    print(\"Trace %d: Textin[0] = %02x  Key[0] = %02x --> SBox Output = %02x --> HW = %d\"%(\n",
    "        n, tin, key, intermediate, hw[intermediate]))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cov(x, y):\n",
    "    # Find the covariance between two 1D lists (x and y).\n",
    "    # Note that var(x) = cov(x, x)\n",
    "    return np.cov(x, y)[0][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we apply that across all the traces in the profiling run. This will let us \"split\" each trace into a different group according to the resulting Hamming weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2: Find HW(sbox) to go with each input\n",
    "# Note - we're only working with ONE byte here\n",
    "target_byte = 0\n",
    "tempSbox = [sbox[project_template.textins[n][target_byte] ^ project_template.keys[n][target_byte]] for n in range(len(project_template.traces))] \n",
    "tempHW   = [hw[s] for s in tempSbox]\n",
    "\n",
    "# 2.5: Sort traces by HW\n",
    "# Make 9 blank lists - one for each Hamming weight\n",
    "tempTracesHW = [[] for _ in range(9)]\n",
    "\n",
    "# Fill them up\n",
    "for i in range(len(project_template.traces)):\n",
    "    HW = tempHW[i]\n",
    "    tempTracesHW[HW].append(project_template.waves[i])\n",
    "\n",
    "# Switch to numpy arrays\n",
    "tempTracesHW = [np.array(tempTracesHW[HW]) for HW in range(9)]\n",
    "\n",
    "#print len(tempTracesHW[8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Points of Interest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After sorting the traces by their Hamming weights, we need to find an \"average trace\" for each weight. We can make an array to hold 9 of these averages:\n",
    "\n",
    "```\n",
    "tempMeans = np.zeros((9, len(tempTraces[0])))\n",
    "```\n",
    "\n",
    "Then, we can fill up each of these traces one-by-one. NumPy's <code>average()</code> function makes this easy by including an \"axis\" input. We can tell NumPy to find the average at each point in time and repeat this for all 9 weights:\n",
    "\n",
    "```\n",
    "for i in range(9):\n",
    "    tempMeans[i] = np.average(tempTracesHW[i], 0)\n",
    "```\n",
    "\n",
    "Once again, it's a good idea to plot one of these averages to make sure that the average traces look okay:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3: Find averages\n",
    "tempMeans = np.zeros((9, len(project_template.waves[0])))\n",
    "for i in range(9):\n",
    "    tempMeans[i] = np.average(tempTracesHW[i], 0)\n",
    "    \n",
    "hv.Curve(tempMeans[2]).opts(height=600, width=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use these average traces to find points of interest using the ''sum of differences'' method. The first step is to create a \"trace\" that stores these differences:\n",
    "\n",
    "```\n",
    "tempSumDiff = np.zeros(len(tempTraces[0]))\n",
    "```\n",
    "\n",
    "Then, we want to look at all of the pairs of traces, subtract them, and add them to the sum of differences. This is simple to do with 2 loops:\n",
    "\n",
    "```\n",
    "for i in range(9):\n",
    "    for j in range(i):\n",
    "        tempSumDiff += np.abs(tempMeans[i] - tempMeans[j])\n",
    "```\n",
    "\n",
    "This sum of differences will have peaks where the average traces showed a lot of variance. Make sure that there are a handful of high peaks that we can use as points of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4: Find sum of differences\n",
    "tempSumDiff = np.zeros(len(project_template.waves[0]))\n",
    "for i in range(9):\n",
    "    for j in range(i):\n",
    "        tempSumDiff += np.abs(tempMeans[i] - tempMeans[j])\n",
    "        \n",
    "hv.Curve(tempSumDiff).opts(height=600, width=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the most interesting points of interest are the highest peaks in this sum of differences plot. However, we can't just sort the array and pick the top points. Remember (because you read the theory page, right? <sup>right?</sup>) that we need to make sure our points have some space between them. For instance, it would be a bad idea to pick both 1950 and 1951 as points of interest.\n",
    "\n",
    "We can use the algorithm from the theory page to pick out some POIs:\n",
    "\n",
    "* Make an empty list of POIs\n",
    "* Find the biggest peak in the sum of differences trace and add it to the list of POIs\n",
    "* Zero out some of the surrounding points\n",
    "* Repeat until we have enough POIs\n",
    "\n",
    "Try to code this on your own - it's a good exercise. If you get stuck, here's our implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5: Find POIs\n",
    "POIs = []\n",
    "numPOIs = 5\n",
    "POIspacing = 5\n",
    "for i in range(numPOIs):\n",
    "    # Find the max\n",
    "    nextPOI = tempSumDiff.argmax()\n",
    "    POIs.append(nextPOI)\n",
    "    \n",
    "    # Make sure we don't pick a nearby value\n",
    "    poiMin = max(0, nextPOI - POIspacing)\n",
    "    poiMax = min(nextPOI + POIspacing, len(tempSumDiff))\n",
    "    for j in range(poiMin, poiMax):\n",
    "        tempSumDiff[j] = 0\n",
    "    \n",
    "#print POIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.Curve(tempSumDiff).opts(height=600, width=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(POIs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating a Template"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With 5 (or <code>numPOIs</code>) POIs picked out, we can build our multivariate distributions at each point for each Hamming weight. We need to write down two matrices for each weight:\n",
    "\n",
    "* A mean matrix (<code>1 x numPOIs</code>) which stores the mean at each POI\n",
    "* A covariance matrix (<code>numPOIs x numPOIs</code>) which stores the variances and covariances between each of the POIs\n",
    "\n",
    "The mean matrix is easy to set up because we've already found the mean at every point. All we need to do is grab the right points:\n",
    "\n",
    "```\n",
    "meanMatrix = np.zeros((9, numPOIs))\n",
    "for HW in range(9):\n",
    "    for i in range(numPOIs):\n",
    "        meanMatrix[HW][i] = tempMeans[HW][POIs[i]]\n",
    "```\n",
    "\n",
    "The covariance matrix is a bit more complex. We need a way to find the covariance between two 1D arrays. Helpfully, NumPy has the <code>cov(a, b)</code> function, which returns the matrix\n",
    "```\n",
    "np.cov(a, b) = [[cov(a, a), cov(a, b)],\n",
    "                [cov(b, a), cov(b, b)]]\n",
    "```\n",
    "\n",
    "We can use this to define our own covariance function:\n",
    "```\n",
    "def cov(x, y):\n",
    "    # Find the covariance between two 1D lists (x and y).\n",
    "    # Note that var(x) = cov(x, x)\n",
    "    return np.cov(x, y)[0][1]\n",
    "```\n",
    "\n",
    "As mentioned in the comments, this function can also calculate the variance of an array by passing the same array in for both <code>x</code> and <code>y</code>.\n",
    "\n",
    "We'll use our function to build the covariance matrices, as below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6: Fill up mean and covariance matrix for each HW\n",
    "meanMatrix = np.zeros((9, numPOIs))\n",
    "covMatrix  = np.zeros((9, numPOIs, numPOIs))\n",
    "for HW in range(9):\n",
    "    for i in range(numPOIs):\n",
    "        # Fill in mean\n",
    "        meanMatrix[HW][i] = tempMeans[HW][POIs[i]]\n",
    "        for j in range(numPOIs):\n",
    "            x = tempTracesHW[HW][:,POIs[i]]\n",
    "            y = tempTracesHW[HW][:,POIs[j]]\n",
    "            covMatrix[HW,i,j] = cov(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally - let's confirm these matricies look A-OK. Basically you need to ensure there isn't zeros in them, which normally indiciates you didn't have enough data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(meanMatrix)\n",
    "print(covMatrix[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Applying the Template"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The very last step is to apply our template to these traces. We want to keep a running total of $\\log P_k = \\sum_j \\log p_{k,j}$, so we'll make space for our 256 guesses:\n",
    "```\n",
    "P_k = np.zeros(256)\n",
    "```\n",
    "\n",
    "Then, we want to do the following for every attack trace:\n",
    "\n",
    "* Grab the samples from the points of interest and store them in a list $\\mathbf{a}$\n",
    "* For all 256 of the subkey guesses...\n",
    " * Figure out which Hamming weight we need, according to our known plaintext and guessed subkey\n",
    " * Build a <code>multivariate_normal</code> object using the relevant mean and covariance matrices\n",
    " * Calculate the log of the PDF ($\\log f(\\mathbf{a})$) and add it to the running total\n",
    "* List the best guesses we've seen so far\n",
    "\n",
    "Make sure you've included the multivariate stats from SciPy too - it's normally an extra import."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "project_validate = cw.open_project(\"projects/tutorial_template_validate.cwp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal\n",
    "# 2: Attack\n",
    "# Running total of log P_k\n",
    "P_k = np.zeros(256)\n",
    "for j in range(len(project_validate.traces)):\n",
    "    # Grab key points and put them in a small matrix\n",
    "    a = [project_validate.waves[j][POIs[i]] for i in range(len(POIs))]\n",
    "    \n",
    "    # Test each key\n",
    "    for k in range(256):\n",
    "        # Find HW coming out of sbox\n",
    "        HW = hw[sbox[project_validate.textins[j][target_byte] ^ k]]\n",
    "    \n",
    "        # Find p_{k,j}\n",
    "        rv = multivariate_normal(meanMatrix[HW], covMatrix[HW])\n",
    "        p_kj = rv.logpdf(a)\n",
    "   \n",
    "        # Add it to running total\n",
    "        P_k[k] += p_kj\n",
    "\n",
    "    # Print our top 5 results so far\n",
    "    # Best match on the right\n",
    "    print(\" \".join([\"%02x\"%j for j in P_k.argsort()[-5:]]))\n",
    "    \n",
    "guess = P_k.argsort()[-1]\n",
    "print(hex(guess))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With any luck you will have found the correct key byte! You can try extending this to attack multiple bytes instead. The traces you already recorded will work just great, you'll need to make the following changes:\n",
    "\n",
    "* Build a new POI for each byte number\n",
    "* Build new matricies for each byte number\n",
    "* Apply those matricies for each byte"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "That's it! This template attack must have raised some new questions for you. A few things to try beyond the above:\n",
    "\n",
    "* Get rid of the Hamming weight assumption. You can build templates based on each key byte value itself directly, which means a SINGLE power trace would result in the secret key(!).\n",
    " * Doing this requires a *lot* more data, since you need to build information on each key byte.\n",
    " * Make your life easier by using a FIXED plaintext. The same plaintext is used during building the template as when you apply the attack."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert guess == 0x2b, \"Failed to break key byte 0, got {}\".format(hex(guess))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
